package com.yanyeori.framework.security.service;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import com.yanyeori.framework.core.constant.BaseConstant;
import com.yanyeori.framework.core.constant.BaseWebCodeEnum;
import com.yanyeori.framework.core.model.WebResponse;
import com.yanyeori.framework.security.bo.InterfaceLimiterBO;
import com.yanyeori.framework.security.cache.PermissionsCache;
import com.yanyeori.framework.security.configprop.SecurityConfigProps;
import com.yanyeori.framework.security.database.InterfaceDao;
import com.yanyeori.framework.security.enums.EnableEnum;
import com.yanyeori.framework.security.util.SecurityUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.InitializingBean;

import javax.annotation.PostConstruct;
import java.lang.ref.SoftReference;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

/**
 * 接口ip限流服务
 * 每秒只发出5个令牌，此处是单进程服务的限流,内部采用令牌捅算法实现
 * RateLimiter.create(n)能获取到n+1个令牌
 * RateLimiter rateLimiter = RateLimiter.create(1);
 *
 * @author chenkuan 2020/12/3
 */
@Slf4j
public class DefaultRateLimiterService implements RateLimiterService, InitializingBean {

    private final SecurityConfigProps config;
    private final InterfaceDao interfaceDao;

    public DefaultRateLimiterService(SecurityConfigProps config, InterfaceDao interfaceDao) {
        this.config = config;
        this.interfaceDao = interfaceDao;
    }

    private static LoadingCache<String, RateLimiter> limiterLoadingCache;

    @Override
    public void afterPropertiesSet() throws Exception {
        limiterLoadingCache = CacheBuilder.newBuilder()
                .maximumSize(config.getLimiter().getCacheMaxiSize())
                .expireAfterAccess(2, TimeUnit.MINUTES)
                .build(new CacheLoader<String, RateLimiter>() {
                    @Override
                    public RateLimiter load(String key) throws Exception {
                        //新的key初始化令牌捅
                        //key=ip|backUrl
                        return RateLimiter.create(PermissionsCache.LIMITER_MAP.get(key).getPermitsPer());
                    }
                });
    }

    /**
     * 校验限流访问服务
     *
     * @param requestIp   请求ip
     * @param servletPath 后端servlet请求路径
     */
    @Override
    public WebResponse<?> validLimiterService(String requestIp, String servletPath) {
        if (!config.getLimiter().isEnable()) {
            return WebResponse.ok();
        }
        InterfaceLimiterBO limiterDto = getLimiterInfo(requestIp, servletPath);
        if (limiterDto != null && limiterDto.getPermitsPer() > 0 && EnableEnum.ENABLE.getCode().equals(limiterDto.getEnabled())) {
            if (!tryAcquire(limiterDto.getKey())) {
                //获取访问令牌失败
                return WebResponse.fail(BaseWebCodeEnum.SERVER_BUSY.code(), BaseWebCodeEnum.SERVER_BUSY.msg());
            }
        }
        return WebResponse.ok();
    }

    /**
     * 获取访问令牌
     */
    private boolean tryAcquire(String key) {
        //获取令牌桶
        RateLimiter rateLimiter = getRateLimiter(key);
        //从令牌桶中取一个令牌
        return rateLimiter != null && rateLimiter.tryAcquire(1, config.getLimiter().getTryAcquireTimeOut(), TimeUnit.MILLISECONDS);
    }

    /**
     * 获取令牌桶
     */
    private RateLimiter getRateLimiter(String key) {
        try {
            return limiterLoadingCache.get(key);
        } catch (ExecutionException e) {
            log.warn("Failed to get token bucket, key={}", key);
            return null;
        }
    }

    /**
     * 获取接口限流配置
     *
     * @return userId|interface|ip
     */
    private InterfaceLimiterBO getLimiterInfo(String requestIp, String servletPath) {
        Map<String, InterfaceLimiterBO> limiterMap;
        if (config.isDevMode()) {
            limiterMap = loadAllLimiterConfig();
        } else {
            limiterMap = PermissionsCache.LIMITER_MAP;
        }
        if (!config.isDevMode()) {
            Map<String, InterfaceLimiterBO> softVerMap = PermissionsCache.VER_LIMITER.get();
            InterfaceLimiterBO limiterDto;
            if (softVerMap != null && (limiterDto = softVerMap.get(getLimiterKey(requestIp, servletPath))) != null) {
                return limiterDto;
            }
        }
        for (InterfaceLimiterBO limiterDto : limiterMap.values()) {
            if (SecurityUtil.matchPath(limiterDto.getBackUrl(), servletPath) && SecurityUtil.matchIp(limiterDto.getIp(), requestIp)) {
                if (config.isDevMode()) {
                    InterfaceLimiterBO oldLimiterDto = PermissionsCache.LIMITER_MAP.get(limiterDto.getKey());
                    PermissionsCache.LIMITER_MAP.put(limiterDto.getKey(), limiterDto);
                    if (oldLimiterDto != null && oldLimiterDto.getPermitsPer().compareTo(limiterDto.getPermitsPer()) != 0) {
                        limiterLoadingCache.refresh(limiterDto.getKey());
                    }
                } else {
                    Map<String, InterfaceLimiterBO> softVerMap = PermissionsCache.VER_LIMITER.get();
                    if (softVerMap == null) {
                        softVerMap = new HashMap<>();
                        PermissionsCache.VER_LIMITER = new SoftReference<>(softVerMap);
                    }
                    softVerMap.put(getLimiterKey(requestIp, servletPath), limiterDto);
                }
                return limiterDto;
            }
        }
        //使用默认全局配置
        double permits = config.getLimiter().getPermitsPerSecond();
        if (permits > 0) {
            InterfaceLimiterBO limiterDto = new InterfaceLimiterBO(servletPath, requestIp, config.getLimiter().getPermitsPerSecond(), EnableEnum.ENABLE.getCode());
            PermissionsCache.LIMITER_MAP.put(limiterDto.getKey(), limiterDto);
            if (!config.isDevMode()) {
                Map<String, InterfaceLimiterBO> softVerMap = PermissionsCache.VER_LIMITER.get();
                if (softVerMap == null) {
                    softVerMap = new HashMap<>();
                    PermissionsCache.VER_LIMITER = new SoftReference<>(softVerMap);
                }
                softVerMap.put(limiterDto.getKey(), limiterDto);
            }
            return limiterDto;
        }
        return null;
    }

    private String getLimiterKey(String ip, String path) {
        return ip + BaseConstant.SIGN_VERTICAL + path;
    }

    /**
     * 查询接口限流配置
     */
    private Map<String, InterfaceLimiterBO> loadAllLimiterConfig() {
        Map<String, InterfaceLimiterBO> map = new HashMap<>();
        List<InterfaceLimiterBO> limiterList = interfaceDao.listLimiters();
        for (InterfaceLimiterBO limiterDto : limiterList) {
            map.putIfAbsent(limiterDto.getKey(), limiterDto);
        }
        return map;
    }

    /**
     * 系统启动时加载所有限流配置缓存
     */
    @PostConstruct
    public void loadLimiterCache() {
        if (config.getLimiter().isEnable() && !config.isDevMode()) {
            flushLimiterCache();
        }
    }

    /**
     * 刷新所有限流配置缓存
     */
    @Override
    public void flushLimiterCache() {
        PermissionsCache.LIMITER_MAP.clear();
        PermissionsCache.LIMITER_MAP.putAll(loadAllLimiterConfig());
    }
}
